package com.jie.component.mongo.core;

import com.jie.component.mongo.base.MongoDocument;
import com.jie.component.mongo.base.MongoRepository;
import com.jie.component.mongo.wrapper.LambdaMongoQueryWrapper;
import com.jie.component.mongo.wrapper.LambdaMongoUpdateWrapper;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageImpl;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;
import org.springframework.data.mongodb.core.BulkOperations;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.aggregation.Aggregation;
import org.springframework.data.mongodb.core.aggregation.AggregationOperation;
import org.springframework.data.mongodb.core.aggregation.TypedAggregation;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

import static org.springframework.data.mongodb.core.query.Criteria.where;

@RequiredArgsConstructor
public abstract class AbstractMongoRepository<T extends MongoDocument> implements MongoRepository<T> {

    protected final MongoTemplate mongoTemplate;

    protected final Class<T> entityClass;

    @Override
    public T save(T entity) {
        if (entity.getCreateTime() == null) {
            entity.setCreateTime(System.currentTimeMillis());
        }
        entity.setUpdateTime(System.currentTimeMillis());
        return mongoTemplate.save(entity);
    }

    @Override
    public List<T> saveBatch(Collection<T> entityList) {
        if (CollectionUtils.isEmpty(entityList)) {
            return List.of();
        }

        long currentTime = System.currentTimeMillis();
        entityList.forEach(entity -> {
            if (entity.getCreateTime() == null) {
                entity.setCreateTime(currentTime);
            }
            entity.setUpdateTime(currentTime);
        });

        return mongoTemplate.insertAll(entityList).stream()
                .map(entity -> (T) entity)
                .collect(Collectors.toList());
    }

    @Override
    public void deleteById(String id) {
        mongoTemplate.remove(Query.query(where("_id").is(id)), entityClass);
    }

    @Override
    public void deleteByIds(Collection<String> ids) {
        if (CollectionUtils.isEmpty(ids)) {
            return;
        }
        mongoTemplate.remove(Query.query(where("_id").in(ids)), entityClass);
    }


    @Override
    public Optional<T> findById(String id) {
        return Optional.ofNullable(mongoTemplate.findById(id, entityClass));
    }

    @Override
    public List<T> findByIds(Collection<String> ids) {
        if (CollectionUtils.isEmpty(ids)) {
            return List.of();
        }
        return mongoTemplate.find(
                Query.query(where("_id").in(ids)),
                entityClass
        );
    }

    @Override
    public List<T> findAll() {
        return mongoTemplate.findAll(entityClass);
    }

    @Override
    public <R> List<R> aggregate(List<AggregationOperation> operations, Class<R> resultClass) {
        if (CollectionUtils.isEmpty(operations)) {
            return List.of();
        }

        TypedAggregation<T> aggregation = Aggregation.newAggregation(entityClass, operations);
        return mongoTemplate.aggregate(aggregation, resultClass)
                .getMappedResults();
    }

    @Override
    public List<T> findByCondition(LambdaMongoQueryWrapper<T> wrapper) {
        Query query = buildQuery(wrapper);
        return mongoTemplate.find(query, entityClass);
    }

    @Override
    public Page<T> findPageByCondition(LambdaMongoQueryWrapper<T> wrapper, Pageable pageable) {
        Query query = buildQuery(wrapper);
        long total = mongoTemplate.count(query, entityClass);
        
        if (total > 0) {
            query.with(pageable);
            List<T> content = mongoTemplate.find(query, entityClass);
            return new PageImpl<>(content, pageable, total);
        }
        
        return new PageImpl<>(List.of(), pageable, 0);
    }

    @Override
    public void update(LambdaMongoUpdateWrapper<T> updateWrapper) {
        Query query = buildQuery(updateWrapper);
        Update update = updateWrapper.getUpdate()
                .set("update_time", System.currentTimeMillis());
        
        mongoTemplate.updateMulti(query, update, entityClass);
    }

    @Override
    public void updateById(String id, LambdaMongoUpdateWrapper<T> updateWrapper) {
        Update update = updateWrapper.getUpdate()
                .set("update_time", System.currentTimeMillis());
        
        mongoTemplate.updateFirst(
                Query.query(where("_id").is(id)),
                update,
                entityClass
        );
    }

    @Override
    public void upsert(LambdaMongoUpdateWrapper<T> updateWrapper) {
        Query query = buildQuery(updateWrapper);
        Update update = updateWrapper.getUpdate()
                .setOnInsert("create_time", System.currentTimeMillis())
                .set("update_time", System.currentTimeMillis());
        
        mongoTemplate.upsert(query, update, entityClass);
    }

    @Override
    public void delete(LambdaMongoQueryWrapper<T> queryWrapper) {
        Query query = buildQuery(queryWrapper);
        mongoTemplate.remove(query, entityClass);
    }

    @Override
    public Optional<T> findOne(LambdaMongoQueryWrapper<T> wrapper) {
        Query query = buildQuery(wrapper);
        return Optional.ofNullable(mongoTemplate.findOne(query, entityClass));
    }

    @Override
    public long count(LambdaMongoQueryWrapper<T> wrapper) {
        Query query = buildQuery(wrapper);
        return mongoTemplate.count(query, entityClass);
    }

    @Override
    public boolean exists(LambdaMongoQueryWrapper<T> wrapper) {
        Query query = buildQuery(wrapper);
        return mongoTemplate.exists(query, entityClass);
    }

    @Override
    public void updateBatch(Collection<String> ids, LambdaMongoUpdateWrapper<T> updateWrapper) {
        if (CollectionUtils.isEmpty(ids)) {
            return;
        }
        
        Query query = Query.query(where("_id").in(ids));
        Update update = updateWrapper.getUpdate()
                .set("update_time", System.currentTimeMillis());
        
        mongoTemplate.updateMulti(query, update, entityClass);
    }

    @Override
    public void bulkUpsert(List<LambdaMongoUpdateWrapper<T>> updateWrappers) {
        if (CollectionUtils.isEmpty(updateWrappers)) {
            return;
        }

        long currentTime = System.currentTimeMillis();
        BulkOperations bulkOps = mongoTemplate.bulkOps(BulkOperations.BulkMode.UNORDERED, entityClass);

        for (LambdaMongoUpdateWrapper<T> wrapper : updateWrappers) {
            Query query = buildQuery(wrapper);
            Update update = wrapper.getUpdate()
                    .setOnInsert("create_time", currentTime)
                    .set("update_time", currentTime);
            bulkOps.upsert(query, update);
        }

        bulkOps.execute();
    }

    @Override
    public <R> Optional<R> aggregateOne(List<AggregationOperation> operations, Class<R> resultClass) {
        if (CollectionUtils.isEmpty(operations)) {
            return Optional.empty();
        }

        TypedAggregation<T> aggregation = Aggregation.newAggregation(entityClass, operations);
        return Optional.ofNullable(mongoTemplate.aggregate(aggregation, resultClass)
                .getUniqueMappedResult());
    }

    @Override
    public long aggregateCount(List<AggregationOperation> operations) {
        if (CollectionUtils.isEmpty(operations)) {
            return 0;
        }

        operations = new ArrayList<>(operations);
        operations.add(Aggregation.count().as("count"));
        
        TypedAggregation<T> aggregation = Aggregation.newAggregation(entityClass, operations);
        return Optional.ofNullable(mongoTemplate.aggregate(aggregation, CountResult.class)
                        .getUniqueMappedResult())
                .map(CountResult::getCount)
                .orElse(0L);
    }

    // 提取公共方法构建Query
    private Query buildQuery(LambdaMongoQueryWrapper<T> wrapper) {
        Query query = new Query();
        if (!CollectionUtils.isEmpty(wrapper.getCriteriaList())) {
            Criteria[] criteriaArray = wrapper.getCriteriaList().toArray(new Criteria[0]);
            query.addCriteria(new Criteria().andOperator(criteriaArray));
        }

        if (!CollectionUtils.isEmpty(wrapper.getOrders())) {
            query.with(Sort.by(wrapper.getOrders()));
        }
        return query;
    }

    @Data
    private static class CountResult {
        private long count;
    }
}